import torch
from torch import nn

class SVM(nn.Module):
    def __init__(self,width,height,intolerance_rate) -> None:
        super().__init__()
        self.width = width
        self.height = height
        self.rate = intolerance_rate
        self.flat = nn.Sequential(
            torch.nn.Flatten(),
        )
        self.w = nn.Linear(width*height,1) #w.weight w.bias
        self.relu = nn.ReLU()

    def forward(self,img,y):
        w2 = 0.5*torch.pow(self.w.weight,2).sum()
        x = self.flat(img)
        x = self.w(x).squeeze(1)
        x = 1-torch.multiply(x,y)
        x = self.relu(x)
        return w2+self.rate*x.sum()
    
    def classify(self,img,label=False):
        x = self.flat(img)
        if label:
            return torch.sign(self.w(x).squeeze(1))
        else:
            return self.w(x).squeeze(1)

class OVA_SVMs(nn.Module):
    def __init__(
            self,width,height,
            intolerance_rate,class_num) -> None:
        super().__init__()
        self.width = width
        self.height = height
        self.rate = intolerance_rate
        self.class_num = class_num
        self.SVM = nn.ModuleList([
            SVM(
                width,height,
                intolerance_rate
            ) for i in range(class_num)
        ])
        self.out = nn.Softmax(1)

    def forward(self,img,y=None):
        loss = 0
        for i in range(self.class_num):
            yi = y
            yi = (yi==i).to(torch.float)*2-1 #==i 1 !=i -1
            loss += self.SVM[i](img,yi)
        return loss

    def classify(self,img):
        score = torch.zeros((img.shape[0],0)).to('cuda')
        for i in range(self.class_num):
            yi = self.SVM[i].classify(img).unsqueeze(1)
            score = torch.hstack((score,yi))
        pred = nn.functional.normalize(score,1,1) 
        return self.out(pred)

